# coding=utf-8
# @author:      ChengJing
# @name:        operate_inp.py
# @datetime:    2021/6/23 0:50
# @software:    PyCharm
# @description:

import numpy as np
import wntr
from scipy.interpolate import interp1d


class OperateModel:
    """
    配置水力模型文件（inp文件）的参数
    config方法配置水力参数，可根据需求进行适当修改
    """

    def __init__(self, wn, old_step=3600, new_step=60):
        """
        Args:
            wn: wntr模块读取的epanet模型
            old_step: 原pattern的时间步长，单位为秒(s)
            new_step: 插值后pattern的时间步长，单位为秒(s)
        """
        self.wn = wn
        self.old_step = old_step
        self.new_step = new_step

    def _config(self):
        """
        配置水力模型的模拟参数
        """
        # 设置水力模型---->压力驱动
        self.wn.options.hydraulic.demand_model = 'PDD'
        self.wn.options.hydraulic.required_pressure = 24
        self.wn.options.hydraulic.minimum_pressure = 0
        self.wn.options.hydraulic.pressure_exponent = 0.5

        # 设置模型模拟时间
        self.wn.options.time.duration = 24 * 3600
        self.wn.options.time.hydraulic_timestep = self.new_step
        self.wn.options.time.pattern_timestep = self.new_step
        self.wn.options.time.report_timestep = self.new_step

    def _dense(self, is_disturb=False):
        """
        对pattern进行插值
        Args:
            is_disturb: bool, 是否需要添加扰动
        """
        for i in self.wn.pattern_name_list:
            pattern = self.wn.get_pattern(i)
            x = np.arange(0, (len(pattern) + 1) * self.old_step, self.old_step)
            y = list(pattern.multipliers)
            y.append(pattern.multipliers[0])
            f = interp1d(x, y, kind='quadratic')
            nx = np.arange(0, len(pattern) * self.old_step, self.new_step)
            ny = f(nx)
            if is_disturb:
                disturb = 1.0 - \
                    np.random.normal(loc=0, scale=0.1, size=ny.shape)
                ny *= disturb
            pattern.multipliers = ny

    def operate_inp(self, file='./inp/model24.inp'):
        """
        Args:
            file: 保存配置修改之后的模型的文件名
        """
        self._config()
        self._dense()
        self.wn.write_inpfile(filename=file, version=2.2)


# todo
class CheckModel:
    """
    说明：
        该类对epanet水力模型的结构误差进行调整，旨在降低管网模型与真实模型之间
        的结构误差。
    """

    def __init__(self, wn):
        pass


if __name__ == '__main__':
    wn = wntr.network.WaterNetworkModel(r'./inp/hydraulic_model.inp')
    OperateModel(wn).operate_inp()
